"""
WebSocket服务 - 实时推送各种业务消息
支持的消息类型：
- verify_result: 实名认证结果
- order: 订单状态变化
- payment: 支付成功/失败
- favorite: 收藏/取消收藏
- review: 评价/回复
- system: 系统消息
- stock: 库存提醒
- coupon: 优惠券通知
"""
from fastapi import WebSocket, WebSocketDisconnect
from typing import Dict, Set
import json
import asyncio
import logging
from utils.redis_utils import get_redis

logger = logging.getLogger(__name__)


class ConnectionManager:
    """WebSocket连接管理器"""
    def __init__(self):
        # user_id -> Set[WebSocket]
        self.active_connections: Dict[int, Set[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, user_id: int):
        """建立连接"""
        await websocket.accept()
        if user_id not in self.active_connections:
            self.active_connections[user_id] = set()
        self.active_connections[user_id].add(websocket)
        logger.info(f"用户 {user_id} WebSocket连接建立，当前连接数: {len(self.active_connections[user_id])}")

    def disconnect(self, websocket: WebSocket, user_id: int):
        """断开连接"""
        if user_id in self.active_connections:
            self.active_connections[user_id].discard(websocket)
            if not self.active_connections[user_id]:
                del self.active_connections[user_id]
        logger.info(f"用户 {user_id} WebSocket连接断开")

    async def send_personal_message(self, message: dict, user_id: int):
        """发送个人消息"""
        if user_id in self.active_connections:
            disconnected = set()
            for connection in self.active_connections[user_id]:
                try:
                    await connection.send_json(message)
                except Exception as e:
                    logger.error(f"发送消息失败: {e}")
                    disconnected.add(connection)
            # 清理断开的连接
            for conn in disconnected:
                self.disconnect(conn, user_id)


# 全局连接管理器
manager = ConnectionManager()


async def websocket_endpoint(websocket: WebSocket, user_id: int):
    """WebSocket端点处理函数"""
    await manager.connect(websocket, user_id)
    redis_client = get_redis()

    try:
        # 发送连接成功消息
        await websocket.send_json({
            "type": "connected",
            "message": "WebSocket连接成功"
        })

        # 持续监听Redis通知队列
        while True:
            # 从Redis读取该用户的通知
            notification_json = redis_client.rpop(f"ws_notifications:{user_id}")
            if notification_json:
                try:
                    notification = json.loads(notification_json)
                    await manager.send_personal_message(notification, user_id)
                    logger.info(f"向用户 {user_id} 推送消息: {notification.get('type')}")
                except Exception as e:
                    logger.error(f"解析或发送通知失败: {e}")

            # 避免CPU占用过高
            await asyncio.sleep(0.5)

    except WebSocketDisconnect:
        manager.disconnect(websocket, user_id)
        logger.info(f"用户 {user_id} WebSocket主动断开")
    except Exception as e:
        logger.error(f"WebSocket错误: {e}", exc_info=True)
        manager.disconnect(websocket, user_id)

